/**
 * @file gather_mm.h
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */

#ifndef MC2_GATHER_MM_H
#define MC2_GATHER_MM_H

#if defined ASCENDC_CPU_DEBUG
#define SET_G_CORE_TYPE_IS_AIV thread_local int g_coreType = 2
#define SET_G_CORE_TYPE_IS_AIC thread_local int g_coreType = 1
#define DTYPE_X1 half
#define DTYPE_Y half
#else
#define SET_G_CORE_TYPE_IS_AIV
#define SET_G_CORE_TYPE_IS_AIC
#endif

#include "kernel_operator_intf.h"
#include "mc2_matmul_compute.h"
#include "all_gather_matmul_custom_tiling.h"

namespace AscendC {
using A_DTYPE = DTYPE_X1;
using B_DTYPE = DTYPE_X1;
using C_DTYPE = DTYPE_Y;
using BIAS_DTYPE = DTYPE_Y;

template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE>
__aicore__ inline void MatmulKernelLocal(GM_ADDR aGM, GM_ADDR bGM, GM_ADDR biasGM, GM_ADDR cGM, AllGatherMatmulRCSTiling &cfg,
    TCubeTiling &tiling, Hccl<HCCL_SERVER_TYPE_AICPU> &hccl)
{
    if ASCEND_IS_AIC {
        if (GetBlockIdx() >= tiling.usedCoreNum) {
            return;
        }
        using C_T = typename C_TYPE::T;
        const auto aRankDataCnt = cfg.rankM * cfg.rankK;
        const auto cRankDataCnt = cfg.rankM * cfg.rankN;

        MatmulCompute<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE> mmLocal;
        mmLocal.Init(cfg, tiling);
        mmLocal.UpdateWeightBias(bGM, biasGM);
        mmLocal.UpdateAddress(aGM, aRankDataCnt, cGM + hccl.GetRankId() * cRankDataCnt * sizeof(C_T), cRankDataCnt);
        mmLocal.Process();
        mmLocal.End();
    }
}

template <class A_TYPE, class B_TYPE, class C_TYPE, class BIAS_TYPE>
__aicore__ inline void MatmulKernel(GM_ADDR aAddr, GM_ADDR bGM, GM_ADDR biasGM, GM_ADDR cAddr, AllGatherMatmulRCSTiling &cfg,
    TCubeTiling &tiling, Hccl<HCCL_SERVER_TYPE_AICPU> &hccl, HcclHandle &handleId, uint32_t tileCnt)
{
    if ASCEND_IS_AIC {
        if (GetBlockIdx() >= tiling.usedCoreNum) {
            for (uint32_t i = 0; i < tileCnt; i++) {
                CrossCoreSetFlag<0x0, PIPE_FIX>(0x8);
                CrossCoreWaitFlag(0x8);
            }
            return;
        }
        using A_T = typename A_TYPE::T;
        using C_T = typename C_TYPE::T;
        const auto aDataCnt = tiling.M * tiling.Ka;
        const auto aOffset = aDataCnt * sizeof(A_T);
        const auto cDataCnt = tiling.M * tiling.N;
        const auto cOffset = cDataCnt * sizeof(C_T);
        const auto aRankOffset = cfg.rankM * cfg.rankK * sizeof(A_T);
        const auto cRankOffset = cfg.rankM * cfg.rankN * sizeof(C_T);

        MatmulCompute<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE> mm;
        mm.Init(cfg, tiling);
        mm.UpdateWeightBias(bGM, biasGM);
        for (uint32_t i = 0; i < tileCnt; i++) {
            // wait current handle allgather
            hccl.Wait(handleId);
            CrossCoreSetFlag<0x0, PIPE_FIX>(0x8);
            CrossCoreWaitFlag(0x8);
            // calculate all ranks data
            for (uint32_t rankId = 0; rankId < hccl.GetRankDim(); rankId++) {
                // skip local rank
                if (rankId == hccl.GetRankId()) {
                    continue;
                }
                mm.UpdateAddress(aAddr + rankId * aRankOffset, aDataCnt, cAddr + rankId * cRankOffset, cDataCnt);
                mm.Process();
            }
            aAddr += aOffset;
            cAddr += cOffset;
        }
        mm.End();
    }
}
}
#endif  // MC2_GATHER_MM_H
